from argparse import ArgumentParser
import time

import zmq


# 服务器状态枚举
# 主机，等待同伴连接
STATE_PRIMARY = 1
# 备机，等待同伴连接
STATE_BACKUP = 2
# 激活态，处理应用程序请求
STATE_ACTIVE = 3
# 被动态，不接收请求
STATE_PASSIVE = 4

# 对话节点事件
# 主机
PEER_PRIMARY = 1
# 备机
PEER_BACKUP = 2
# 激活态
PEER_ACTIVE = 3
# 被动态
PEER_PASSIVE = 4
# 客户端请求
CLINENT_REQUEST = 5

# 发送状态信息的间隔时间
# 如果对方在两次心跳过后都没有应答，则视为断开
HEARTBEAT = 1000


# 有限状态机
class BStarsState(object):
    def __init__(self, state, event, peer_expiry):
        # 当前状态
        self.state = state
        # 当前事件
        self.event = event
        # 判定节点死亡的时限
        self.peer_expiry = peer_expiry


class BStarsException(Exception):
    pass


fsm_states = {
    # 主机等待同伴连接 该状态下接收CLIENT_REQUEST事件
    STATE_PRIMARY: {
        PEER_BACKUP: ("I: 已连接至备机（slave），可以作为master运行", STATE_ACTIVE),
        PEER_ACTIVE: ("I: 已连接至备机（master），可以作为slave运行", STATE_PASSIVE)
    },
    # 备机等待同伴连接 该状态下拒绝CLIENT_REQUEST事件
    STATE_BACKUP: {
        PEER_ACTIVE: ("I: 已连接至主机（master），可以作为slave运行", STATE_PASSIVE),
        CLINENT_REQUEST: ("", False)
    },
    # 服务器处于激活态 该状态下接受CLIENT_REQUEST事件
    STATE_ACTIVE: {
        PEER_ACTIVE: ("E: 严重错误：双master。正在退出", False)
    },
    # 服务器处于被动态 若同伴已死，CLIENT_REQUEST事件将触发故障恢复
    STATE_PASSIVE: {
        PEER_PRIMARY: ("I: 主机（slave）正在重启，可作为master运行", STATE_ACTIVE),
        PEER_BACKUP: ("I: 备机（slave）正在重启，可作为master运行", STATE_ACTIVE),
        PEER_PASSIVE: ("E: 严重错误：双slave。正在退出", False),
        CLINENT_REQUEST: (CLINENT_REQUEST, True)
    }
}


# 执行有限状态机（将事件绑定至状态）
def run_fsm(fsm):
    state_dict = fsm_states.get(fsm.state, {})
    res = state_dict.get(fsm.event)
    if res:
        msg, state = res
    else:
        return
    if state is False:
        raise BStarsException(msg)
    elif msg == CLINENT_REQUEST:
        # 若心跳超时，同伴将成为master 此行为由客户端请求触发。
        assert fsm.peer_expiry > 0
        if int(time.time() * 1000) > fsm.peer_expiry:
            # 同伴已死，转为激活态
            fsm.state = STATE_ACTIVE
        else:
            # 同伴还在，拒绝请求
            raise BStarsException()
    else:
        print('msg: ', msg)
        fsm.state = state
        print("set fsm.state: ", state)


def main():
    parser = ArgumentParser()
    group = parser.add_mutually_exclusive_group()
    group.add_argument("-p", "--primary", action="store_true", default=False)
    group.add_argument("-b", "--backup", action="store_true", default=False)
    args = parser.parse_args()

    ctx = zmq.Context()
    statepub = ctx.socket(zmq.PUB)
    statesub = ctx.socket(zmq.SUB)
    statesub.setsockopt_string(zmq.SUBSCRIBE, "")
    frontend = ctx.socket(zmq.ROUTER)

    fsm = BStarsState(0, 0, 0)

    if args.primary:
        print("I: Primary master, waiting for backup (slave)")
        frontend.bind("tcp://*:5001")
        statepub.bind("tcp://*:5003")
        statesub.connect("tcp://localhost:5004")
        print('set fsm.state: 1')
        fsm.state = STATE_PRIMARY
    elif args.backup:
        print("I: Backup slave, waiting for primary (master)")
        frontend.bind("tcp://*:5002")
        statepub.bind("tcp://*:5004")
        statesub.connect("tcp://localhost:5003")
        print('set fsm.state: 2')
        fsm.state = STATE_BACKUP

    send_state_at = int(time.time() * 1000 + HEARTBEAT)
    poller = zmq.Poller()
    poller.register(frontend, zmq.POLLIN)
    poller.register(statesub, zmq.POLLIN)

    while True:
        time_left = send_state_at - int(time.time() * 1000)
        if time_left < 0:
            time_left = 0
        socks = dict(poller.poll(time_left))
        if socks.get(frontend) == zmq.POLLIN:
            msg = frontend.recv_multipart()
            fsm.event = CLINENT_REQUEST
            try:
                run_fsm(fsm)
                frontend.send_multipart(msg)
            except BStarsException:
                del msg
        if socks.get(statesub) == zmq.POLLIN:
            msg = statesub.recv()
            fsm.event = int(msg)
            if args.primary:
                print("primary statesub recv ", fsm.event)
            elif args.backup:
                print("backup statesub recv ", fsm.event)
            del msg
            try:
                run_fsm(fsm)
                fsm.peer_expiry = int(time.time() * 1000) + (2 * HEARTBEAT)
            except BStarsException:
                break
        if int(time.time() * 1000) >= send_state_at:
            if args.primary:
                print("primary statepub send ", fsm.state)
            elif args.backup:
                print("backup statepub send ", fsm.state)
            statepub.send_string("%d" % fsm.state)
            send_state_at = int(time.time() * 1000) + HEARTBEAT


if __name__ == "__main__":
    main()
